#!/usr/bin/env python
# -*- coding:utf-8 -*-
# @FileName  :kan_conv_mnist.py
# @Time      :2024/10/23 15:27
# @Author    :Jasonljl


#!/usr/bin/env python
# -*- coding:utf-8 -*-
# @FileName  :cnn_mnist.py
# @Time      :2024/10/23 15:11
# @Author    :Jasonljl


#!/usr/bin/env python
# -*- coding:utf-8 -*-
# @FileName  :resnet_mnist.py
# @Time      :2024/10/23 11:28
# @Author    :Jasonljl


#!/usr/bin/env python
# -*- coding:utf-8 -*-
# @FileName  :mnist.py
# @Time      :2024/10/23 10:50
# @Author    :Jasonljl


import torch
import torch.nn as nn
import torchvision.datasets
from torchvision.datasets import mnist,SBDataset
from torchvision.models import resnet18
import torch.optim as optim
import torchvision.transforms as transforms
from torch.utils.data import DataLoader
from tqdm import tqdm
import matplotlib.pyplot as plt
import torch.nn.functional as F
#!/usr/bin/env python
# -*- coding:utf-8 -*-
# @FileName  :cnn_mnist.py
# @Time      :2024/10/23 15:11
# @Author    :Jasonljl


#!/usr/bin/env python
# -*- coding:utf-8 -*-
# @FileName  :resnet_mnist.py
# @Time      :2024/10/23 11:28
# @Author    :Jasonljl


#!/usr/bin/env python
# -*- coding:utf-8 -*-
# @FileName  :mnist.py
# @Time      :2024/10/23 10:50
# @Author    :Jasonljl


import torch
import torch.nn as nn
import torchvision.datasets
from torchvision.datasets import mnist,SBDataset
from torchvision.models import resnet18
import torch.optim as optim
import torchvision.transforms as transforms
from torch.utils.data import DataLoader
from tqdm import tqdm
import matplotlib.pyplot as plt
import torch.nn.functional as F
import torch
from torch import nn
import torch.nn.functional as F
import sys
sys.path.append(r"H:\python\KAN\KAN_VAE_MNIST-main\KAN_VAE_MNIST-main\Convolutional-KANs-master\Convolutional-KANs-master")
from kan_convolutional.KANConv import KAN_Convolutional_Layer


class KANC_MLP(nn.Module):
    def __init__(self, device: str = 'cpu'):
        super().__init__()
        self.conv1 = KAN_Convolutional_Layer(
            n_convs=5,
            kernel_size=(3, 3),
            device=device
        )

        self.conv2 = KAN_Convolutional_Layer(
            n_convs=5,
            kernel_size=(3, 3),
            device=device
        )

        self.pool1 = nn.MaxPool2d(
            kernel_size=(2, 2)
        )

        self.flat = nn.Flatten()

        self.linear1 = nn.Linear(625, 256)
        self.linear2 = nn.Linear(256, 10)

    def forward(self, x):
        x = self.conv1(x)

        x = self.pool1(x)

        x = self.conv2(x)
        x = self.pool1(x)
        x = self.flat(x)
        x = self.linear1(x)
        x = self.linear2(x)
        x = F.log_softmax(x, dim=1)
        return x
transforms=transforms.Compose([
    transforms.ToTensor(),
    transforms.Resize((28,28)),
    transforms.Normalize((0.5,),(0.5))
])

trainsets=torchvision.datasets.MNIST(
    root="./mnist",download=False,train=True,transform=transforms
)
testsets=torchvision.datasets.MNIST(
    root="./mnist",download=False,train=False,transform=transforms
)

trainloaders=DataLoader(dataset=trainsets,batch_size=64,shuffle=True)
testloaders=DataLoader(dataset=testsets,batch_size=64,shuffle=False)

model=KANC_MLP()
device=torch.device("cuda" if torch.cuda.is_available() else "cpu")

model.to(device)
optimizer=optim.Adam(model.parameters(),lr=1e-3,weight_decay=1e-4)
criterion=nn.CrossEntropyLoss()
acc_list = []
val_list = []
for epoch in range(10):
    model.train()
    with tqdm(trainloaders) as phar:
        for j,(images,labels) in enumerate(phar):
            # print(images.shape)
            images=images.to(device)
            # images=images.repeat(1,3,1,1)
            # print(images.shape)
            optimizer.zero_grad()
            output=model(images)
            loss=criterion(output,labels)
            loss.backward()
            optimizer.step()
            acc=(output.argmax(dim=1)==labels.to(device)).float().mean()
            phar.set_postfix(loss=loss.item(),acc=acc.item(),lr=optimizer.param_groups[0]["lr"])
    model.eval()
    val_loss=0
    acc=0
    with torch.no_grad():
        for images,labels in testloaders:
            images = images.to(device)
            # images=images.repeat(1,3,1,1)
            output=model(images)
            val_loss+=criterion(output,labels.to(device)).item()
            acc+=(
                (output.argmax(dim=1)==labels.to(device)).float().mean().item()
            )
    val_loss/=len(testloaders)
    acc/=len(testloaders)
    acc_list.append(acc)
    val_list.append(val_loss)

    print(f"kan_cnn_epoch:{epoch + 1},val loss:{val_loss},val_acc:{acc}\n")
    with open("result_log.txt", "a+") as f:
        res = f"kan_cnn_epoch:{epoch + 1},val loss:{val_loss},val_acc:{acc}\n"
        f.write(res)
print("done!!!")

epochs=range(1,11)
plt.plot(epochs,acc_list,label="Validation Accuracy")
plt.xlabel("Epoch")
plt.ylabel("Accuracy")
plt.title("Validation Accuracy per Epoch")
plt.legend()
plt.show()
plt.savefig("./kan_mnist.png")
# print(f"epoch:{epoch+1},val loss:{val_loss},val_acc:{acc}")
